package org.example.aigc.utils;


import java.io.*;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.channels.FileChannel;
import java.nio.file.Files;
import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;

/**
 * 处理 wav 文件的工具类
 */
public class WavHandle {

    /**
     * 字节流数据
     */
    private final List<byte[]> wavFiles;

    private ByteBuffer wavHeader;

    public WavHandle() {
        this.wavFiles = new ArrayList<>();
        this.wavHeader = null;
    }

    public WavHandle addWavFile(Path wavFilePath) throws IOException {
        byte[] wavData = Files.readAllBytes(wavFilePath);
        this.wavFiles.add(wavData);

        // If this is the first file, keep its header
        if (this.wavHeader == null) {
            this.wavHeader = ByteBuffer.wrap(wavData, 0, 44);
        }
        return this;
    }

    public void build(Path outputFilePath) throws IOException {
        if (wavFiles.isEmpty()) {
            throw new IllegalStateException("No WAV files added to the factory.");
        }

        // Calculate the total data size
        int totalDataSize = 0;
        for (byte[] data : wavFiles) {
            totalDataSize += data.length - 44; // Exclude the header
        }

        // Update the data chunk size in the header
        wavHeader.position(40);
        wavHeader.putInt(totalDataSize);

        // Write the merged file
        try (OutputStream outputStream = new FileOutputStream(outputFilePath.toFile())) {
            // Write the header
            outputStream.write(wavHeader.array());

            // Write the audio data from all WAV files
            wavFiles.remove(0);
            for (byte[] data : wavFiles) {
                outputStream.write(data, 44, data.length - 44); // Write the audio data, excluding the header
            }
        }
    }

    /**
     * 获取时长
     *
     * @param wavFile 音频/视频文件
     * @return 返回音频时长
     */
    public static double getWavLengthInSeconds(File wavFile) throws IOException {
        FileInputStream fileInputStream = new FileInputStream(wavFile);
        FileChannel fileChannel = fileInputStream.getChannel();
        // WAV文件头部大小，通常为44字节（对于PCM编码的WAV文件）
        ByteBuffer header = ByteBuffer.allocate(44);
        header.order(ByteOrder.LITTLE_ENDIAN);

        // 读取WAV文件头部
        if (fileChannel.read(header) != 44) {
            throw new IOException("Not a valid WAV file");
        }

        // 获取采样率（每秒样本数），通常位于第24-27字节
        int sampleRate = header.getInt(24);

        // 获取数据块大小（字节），通常位于第40-43字节
        int dataChunkSize = header.getInt(40);

        // 获取位深度（每个样本的字节数），通常位于第34-35字节，然后除以8转换为字节
        short bitDepth = header.getShort(34);
        int bytesPerSample = bitDepth / 8;

        // 计算WAV文件的总样本数
        long totalSamples = dataChunkSize / bytesPerSample;

        fileChannel.close();
        fileInputStream.close();
        // 计算WAV文件的长度（秒）
        return (double) totalSamples / sampleRate;
    }
}
